# -*- coding: utf-8 -*-
"""
@Time    : 2024/8/8 10:53 
@Author  : ZhangShenao 
@File    : 1.多向量索引.py
@Desc    : 使用MultiVectorRetriever实现多向量索引

对于一条文档,如果可以从多个维度对文档进行Embedding,记录多维度的特征,则会大大增加相似文档被检索到的概率
常用的多向量(多特征)生成方式包括:
1. 文档切割: 把文档切割成更小的块,通过检索更小的块,查找其父类文档(ParentDocumentRetriever)
2. 摘要: 使用LLM为每个文档块生成一段摘要,将其和原文档一起Embedding,通过检索摘要,返回原文档
3. 假设性问题: 使用LLM为每个文档块生成适合回答的假设性问题,将其和原文档一起Embedding,通过检索假设性问题,返回原文档

通过这种方式可以为一个文档块生成多条特征/向量,在检索时能提升关联文档被检索到的概率

LangChain中封装了MultiVectorRetriever,可以方便的进行多向量检索
"""
import os
import time
import uuid
from typing import List

import dotenv
from langchain.retrievers import MultiVectorRetriever
from langchain.storage import LocalFileStore
from langchain_community.document_loaders import UnstructuredFileLoader
from langchain_community.embeddings import ZhipuAIEmbeddings
from langchain_core.documents import Document
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langchain_pinecone import PineconeVectorStore
from langchain_text_splitters import RecursiveCharacterTextSplitter

dotenv.load_dotenv()

# 加载并切割原始文档
loader = UnstructuredFileLoader("./docs/电商产品数据.txt")
splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
docs = loader.load_and_split(splitter)

# 生成原始文档的摘要信息
llm = ChatOpenAI(model_name="gpt-4o-mini", temperature=0, openai_api_base=os.getenv("OPENAI_API_BASE"))
summarize_chain = (
        {"doc": lambda x: x.page_content}
        | ChatPromptTemplate.from_template("请生成下面文本的摘要：\n\n{doc}")
        | llm
        | StrOutputParser()
)
summaries = summarize_chain.batch(inputs=docs, config={
    "max_concurrency": 1,
})

print("----------生成摘要成功-------------")
time.sleep(1)

# 生成元素文档的唯一doc_id
doc_ids = [str(uuid.uuid4()) for _ in range(len(docs))]

# 创建摘要文档,并在摘要文档的meta_data中记录原始文档的doc_id
summarize_docs: List[Document] = []
for idx, summary in enumerate(summaries):
    doc_id = doc_ids[idx]
    summary_doc = Document(page_content=summary, metadata={"doc_id": doc_id})
    summarize_docs.append(summary_doc)

# 将摘要文档保存至向量数据库
embeddings = ZhipuAIEmbeddings(model="embedding-3")
vector_store = PineconeVectorStore(
    index_name="rag",  # 指定索引名称
    embedding=embeddings,  # 指定Embedding模型
    namespace="llm-ops",  # 指定Namespace
    text_key="text",  # 指定文本属性key
)
vector_store.add_documents(summarize_docs)
print("----------保存摘要文档成功-------------")
time.sleep(5)

# 构造MultiVectorRetriever多向量检索器
file_store = LocalFileStore(root_path="./file_store/")
retriever = MultiVectorRetriever(
    vectorstore=vector_store,  # 指定向量数据库,用于相似性检索
    byte_store=file_store,  # 指定文件数据库,用于根据摘要中的doc_id,获取原始文档
    id_key="doc_id",  # 指定原始文档ID属性key
)

# 将原始文档存储至文件数据库
retriever.docstore.mset(key_value_pairs=list(zip(doc_ids, docs)))
print("----------保存原始文档成功-------------")
time.sleep(5)

# 执行多向量检索
# 通过摘要文档进行相似性检索,并通过doc_id查询原始文档,最终返回原始文档
result = retriever.invoke("请帮我推荐一些潮州特产")
print(f"检索结果: \n{result}")
